import torch


def demonstrate_complete_flow():
    # 模拟模型输出（批量处理，2个样本）
    outputs = torch.tensor([
        [2.1, 0.5, -1.2, 1.8, 0.3, -0.5, 0.9, 1.1, -0.8, 0.2],  # 样本1
        [0.1, 3.2, 0.5, -0.8, 1.2, 0.7, -1.1, 0.9, 0.3, 2.1]  # 样本2
    ])
    print(f"批量输出形状: {outputs.shape}")  # torch.Size([2, 10])

    # 1. 转换为概率
    probabilities = torch.softmax(outputs, dim=1)
    print(f"\n概率分布:")
    for i in range(2):
        print(f"样本{i}: {[f'{p:.3f}' for p in probabilities[i]]}")

    # 2. 获取预测类别
    predicted_classes = torch.argmax(probabilities, dim=1)
    print(f"\n预测类别: {predicted_classes}")  # tensor([0, 1])

    # 3. 获取每个样本的置信度
    confidences = []
    for i in range(len(probabilities)):
        conf = probabilities[i][predicted_classes[i]].item()
        confidences.append(conf)

    print(f"\n预测结果:")
    for i in range(2):
        print(f"样本{i}: 预测类别={predicted_classes[i].item()}, 置信度={confidences[i]:.4f}")


# 运行示例
demonstrate_complete_flow()
